#pragma once

#include "AssocStateBase.h"

namespace AsyncLib
{

template<typename T>
class CAssocState : public CAssocStateBase
{
public:
	template<typename... Args>
	void SetValue(Args&&... args)
	{
		std::unique_lock lock(m_mutex);
		assert(m_status == Pending);

		if (m_status == Pending)
		{
			m_status = Completed;
			m_value.emplace(std::forward<Args>(args)...);
		}
	}

	decltype(auto) GetValue()
	{
		std::unique_lock lock(m_mutex);
		RethrowException();
		assert(m_status == Completed && !m_bRetrieved);

		if (m_status == Completed && !m_bRetrieved)
		{
			m_bRetrieved = !(std::is_reference_v<T> || std::is_trivially_copyable_v<T>);
			return std::forward<T>(m_value.value());
		}

		throw std::logic_error("No value");
	}

private:
	using WrapT = std::conditional_t<std::is_reference_v<T>, std::reference_wrapper<std::remove_reference_t<T>>, T>;
	std::optional<WrapT> m_value;
	bool m_bRetrieved = false;
};

// Suspend the coroutine and resume it in another thread
template<>
class CAssocState<nullptr_t> : public CAssocStateBase
{
public:
	void SuspendFirst(bool bAllowCapture)
	{
		// Make sure to suspend the coroutine first
		m_bCanSetValue = false;
		m_bAllowCapture = bAllowCapture;
	}

	Status SetCallback(std::function<void(bool)>&& fnCallback)
	{
		auto status = CAssocStateBase::SetCallback(std::move(fnCallback));

		std::unique_lock lock(m_mutex);
		m_bCanSetValue = true;
		m_condVar.notify_one();

		return status;
	}

	void SetValue(nullptr_t)
	{
		std::unique_lock lock(m_mutex);

		while (!m_bCanSetValue)
		{
			m_condVar.wait(lock);
		}

		assert(m_status == Pending);

		if (m_status == Pending)
			m_status = Completed;
	}

	nullptr_t GetValue()
	{
		std::unique_lock lock(m_mutex);
		RethrowException();
		return nullptr;
	}

	void CaptureThread()
	{
		// CaptureThread is not allowed for SwitchToTaskThread and SwitchToThreadPool
		assert(m_bAllowCapture);

		if (m_bAllowCapture)
			CAssocStateBase::CaptureThread();
	}

private:
	std::condition_variable m_condVar;
	bool m_bCanSetValue = true;
	bool m_bAllowCapture = true;
};

template<>
class CAssocState<void> : public CAssocStateBase
{
public:
	void SetValue()
	{
		std::unique_lock lock(m_mutex);
		assert(m_status == Pending);

		if (m_status == Pending)
			m_status = Completed;
	}

	void GetValue()
	{
		std::unique_lock lock(m_mutex);
		RethrowException();
	}
};

template<typename T, size_t N>
class CAssocStateN : public CAssocState<T>
{
public:
	void *GetStorage()
	{
		return &m_storage;
	}

private:
	union
	{
		max_align_t dummy1;
		char dummy2[N];
	} m_storage{};
};

} // namespace AsyncLib
